package com.edata.bigdata.algorithm.networks

import com.edata.bigdata.annotations.Edata_Algorithm
import com.edata.bigdata.basic.Algorithm
import org.apache.spark.graphx.{EdgeTriplet, Graph, Pregel, VertexId}
import org.apache.spark.sql.SparkSession
import spire.ClassTag


@Edata_Algorithm(target = "KCORES")
class KCores extends Algorithm with Pregel[Int] {
  override var SESSION: SparkSession = _
  var core_num: Int = 0

  override def mergeMsgs(msg1: Int, msg2: Int): Int = {
    /*
    * msg1 + msg2是最终接受消息的节点需要减掉的度，因为有响应的相邻节点已经去掉了
    * */
    if (msg1 < 0 || msg2 < 0) {
      -1
    } else {
      msg1 + msg2
    }
  }

  override def sendMsg(edge: EdgeTriplet[Int, _]): Iterator[(VertexId, Int)] = {
    if (edge.srcAttr < 0 || edge.dstAttr < 0) {
      /*如果当前A->B元组的A属性或者B属性都小于0，则不发生消息传递*/
      Iterator.empty
    } else if (edge.srcAttr < core_num && edge.dstAttr < core_num) {
      /*如果当前A->B元组的A属性和B属性都小于core_num，则两个节点都不符合要求*/
      Iterator((edge.srcId, -1), (edge.dstId, -1))
    } else if (edge.srcAttr < core_num) {
      /*如果当前A->B元组的A属性小于0，B属性大于0，则A不符合要求，由于A去掉了，则B的度需要减1*/
      Iterator((edge.srcId, -1), (edge.dstId, 1))
    } else if (edge.dstAttr < core_num) {
      /*如果当前A->B元组的B属性小于0，A属性大于0，则B不符合要求，由于B去掉了，则A的度需要减1*/
      Iterator((edge.srcId, 1), (edge.dstId, -1))
    } else {
      /*如果当前A->B元组的A，B属性都大于，则不发送消息*/
      Iterator.empty
    }
  }

  override def handleMsg(vId: VertexId, vAttr: Int, msg: Int): Int = {
    if (msg < 0) {
      /*不符合度大于等于core_num的节点，属性都会直接置为-1*/
      -1
    } else {
      /*去掉不符合要求的相邻节点后，节点剩下的度，依然满足大于等于core_num的要求*/
      math.max(vAttr - msg, 0)
    }
  }

  def run[VD: ClassTag, ED: ClassTag](G: Graph[VD, ED], k: Int): Graph[VD, ED] = {
    core_num = k
    /*
    *
    * G.degrees 本质上返回一个VertexRDD，每个节点附带了'度'属性
    * @outerJoinVertices：G 和 'G.degrees' 首先通过 outerJoinVertices 形成（vid,left_attr,right_attr）的数据元组，
    * 再通过mapfunc对该元组进行进一步处理。
    *
    * */
    val KCGraph = G.outerJoinVertices(G.degrees)((vid, attr, newData) => newData.getOrElse(0))
    val msg = 0
    val PGraph = Pregel(KCGraph, msg)(vprog = handleMsg, sendMsg = sendMsg, mergeMsg = mergeMsgs)
    val SubGraph = PGraph.subgraph(vpred = (vid, v) => v > 0)
    G.mask(SubGraph)
  }
}
